

import logging
import os
import re
import sys
from typing import TypedDict

import numpy as np
import pandas as pd

from SimpleOutcar import SimpleOutcar
from force_view import plot_force_field

VERSION = "0.2.0"
INFO = f"""
------------------------------------------
Distortion Calculator (ver{VERSION})
------------------------------------------
功能描述:
- 从VASP计算的OUTCAR文件中提取原子位置、晶格参数、能量数据、力场信息
- 生成力场可视化图像和力场数据文件
- 计算结构畸变参数(平均位移、相对畸变量)
- 自动生成Excel报告(包含原始坐标、位移向量和计算结果)

开发信息:
- 开发团队: FZU-MCMF
- 主要开发者: Gao Min-Liang, Fang Xuan, Wu Bo*
- 联系方式:
  - 邮箱: wubo@fzu.edu.cn
  - QQ: 654489521

使用说明:
1. 程序会自动搜索目标文件夹下的r7/r3系列计算结果
2. 程序会在每个OUTCAR目录下生成力场信息:
  - force_field.png: 力场可视化图像
  - force_field.csv: 力场数据文件
3. 自动生成包含以下工作表的Excel文件:
  - result: 核心计算结果
  - r3-r7: 原子位移详细数据
------------------------------------------
"""

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")


class DistortionData(TypedDict):
    name: str
    nn_dist: float
    disto_total: float
    disto_avg: float
    disto_rel: float
    e0_total: float
    e0_avg_ev: float
    e0_avg_kj: float


def get_outcar_r7r3(folder: str):
    subfolders = [f for f in os.listdir(folder) if os.path.isdir(os.path.join(folder, f))]

    r7_outcar = ""
    r3_outcar = ""
    subfolders.sort(reverse=True)
    for subfolder in subfolders:
        if match := re.match(r"[rR]7\d?\Z", subfolder):
            outcar = os.path.abspath(os.path.join(folder, subfolder, "OUTCAR"))
            if os.path.exists(outcar):
                r7_outcar = outcar
                break
    for subfolder in subfolders:
        if match := re.match(r".*[rR]3\d?\Z", subfolder):
            outcar = os.path.abspath(os.path.join(folder, subfolder, "OUTCAR"))
            if os.path.exists(outcar):
                r3_outcar = outcar
                break

    return r7_outcar, r3_outcar


def calc_distortion(name, r7: SimpleOutcar, r3: SimpleOutcar):
    atom_num = len(r7.atoms)

    basis_length_r3 = np.linalg.norm(r3.atoms.cell, axis=1)
    logging.info(f"basis_length_r3: {basis_length_r3}")

    # pos_r7 = r7.atoms.cartesian_coords
    # pos_r3 = r3.atoms.cartesian_coords
    # pos_delta = pos_r3 - pos_r7
    pos_delta = r3.positions - r7.positions
    for i in range(pos_delta.shape[0]):
        delta = pos_delta[i, :]
        for j, v in enumerate(delta - basis_length_r3 / 2.0):
            if v > 0.01:
                pos_delta[i, j] -= basis_length_r3[j]

    distos = np.linalg.norm(pos_delta, axis=1)
    disto_total = np.sum(distos)
    disto_avg = disto_total / (atom_num - 1)
    disto_rel = disto_avg / r7.nn_dist
    logging.info(f"disto_total:{disto_total}, disto_avg:{disto_avg}, disto_rel:{disto_rel}")

    e0_total = r3.e0 - r7.e0
    e0_avg_ev = e0_total / (atom_num - 1)
    e0_avg_kj = e0_avg_ev * 96.485
    logging.info(f"e0_total:{e0_total}, e0_avg_ev:{e0_avg_ev}, e0_avg_kj:{e0_avg_kj}")

    res = DistortionData(name=name,
                         nn_dist=r7.nn_dist,
                         disto_total=disto_total,
                         disto_avg=disto_avg,
                         disto_rel=disto_rel,
                         e0_total=e0_total,
                         e0_avg_ev=e0_avg_ev,
                         e0_avg_kj=e0_avg_kj)
    return pos_delta, res


def get_folder(folder: str):
    while True:
        folder = folder or input("Enter the folderpath >>> ").strip()
        if os.path.isdir(folder):
            return os.path.abspath(folder)
        else:
            logging.warning(f"{folder} not found. Please try again.")
            folder = ""


def main():
    folder = sys.argv[1] if len(sys.argv) > 1 else ""
    while True:
        print(INFO)
        try:
            folder = get_folder(folder)
            r7_outcar, r3_outcar = get_outcar_r7r3(folder)

            if not (r7_outcar and r3_outcar):
                logging.error(f"OUTCAR not found in {folder}")
                folder = ""
                continue

            # Plot force field
            plot_force_field(r7_outcar)
            plot_force_field(r3_outcar)
            
            name = os.path.basename(folder)
            r7 = SimpleOutcar(r7_outcar)
            r3 = SimpleOutcar(r3_outcar)
            delta, data = calc_distortion(name, r7, r3)

            # DataFrame
            res_df = pd.DataFrame([data])
            r7_df = pd.DataFrame(r7.atoms.cartesian_coords, columns=["r7_x", "r7_y", "r7_z"])
            r3_df = pd.DataFrame(r3.atoms.cartesian_coords, columns=["r3_x", "r3_y", "r3_z"])
            delta_df = pd.DataFrame(delta, columns=["dx", "dy", "dz"])
            delta_df = pd.concat([r7_df, r3_df, delta_df], axis=1)

            output = os.path.join(folder, "distortion.xlsx")
            with pd.ExcelWriter(output) as writer:
                res_df.to_excel(writer, sheet_name="result", index=False)
                # r7_df.to_excel(writer, sheet_name="r7", index=False)
                # r3_df.to_excel(writer, sheet_name="r3", index=False)
                delta_df.to_excel(writer, sheet_name="r3-r7", index=False)

            logging.info(f"Result saved to {output}")
            folder = ""

        except KeyboardInterrupt:
            print("\n")
            logging.info("Thank you for using Distortion Calculator! Exiting...")
            sys.exit(0)
        except Exception as e:
            print("\n")
            logging.exception(e)
            input("\nPress Enter to Exit >>> ")
            sys.exit(1)


if __name__ == "__main__":
    main()
